package com.simple.llm.service.impl;

import cn.hutool.core.util.IdUtil;
import com.github.zuihou.base.R;
import com.simple.llm.config.LLMProperties;
import com.simple.llm.domain.dto.LLMStreamResultDTO;
import com.simple.llm.domain.dto.NormalChatDTO;
import com.simple.llm.domain.dto.StreamChatDTO;
import com.simple.llm.domain.vo.LLMResponseVO;
import com.simple.llm.service.LLMStrategy;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.ChatCompletionRequest;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ModelApiResponse;
import com.zhipu.oapi.service.v4.model.Usage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.servlet.http.HttpServletResponse;
import java.util.List;
import java.util.function.Consumer;
import java.util.stream.Collectors;

@Slf4j
public class GLMLLMStrategyImpl implements LLMStrategy {
    private final LLMProperties.GptProperties gptProperties;
    private final ClientV4 client;

    public GLMLLMStrategyImpl(LLMProperties.GptProperties gptProperties) {
        this.gptProperties = gptProperties;
        client = new ClientV4.Builder(gptProperties.getApiKeys().get(0)).build();
    }

    @Override
    public R<LLMResponseVO> chat(NormalChatDTO normalChatDTO) {
        List<ChatMessage> messages = normalChatDTO.getMessages().stream().map(o -> new ChatMessage(o.getRole(), o.getContent(), null, null, null)).collect(Collectors.toList());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM3TURBO)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(IdUtil.fastUUID())
                .maxTokens(gptProperties.getMaxContextTokens())
                .temperature(normalChatDTO.getTemperature() == null ? 0.8F : normalChatDTO.getTemperature())
                .topP(normalChatDTO.getTopP() == null ? 0.75F : normalChatDTO.getTopP())
                .build();
        try {
            ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
            if(invokeModelApiResp.getCode() == 200){
                Usage usage = invokeModelApiResp.getData().getUsage();
                String content = invokeModelApiResp.getData().getChoices().get(0).getMessage().getContent().toString();
                LLMResponseVO vo = LLMResponseVO.builder()
                        .content(content)
                        .inputTokenCount((long)usage.getPromptTokens())
                        .outputTokenCount((long)usage.getCompletionTokens())
                        .desc(normalChatDTO.getDesc())
                        .createUser(normalChatDTO.getCreateUser())
                        .build();
                return R.success(vo);
            }
        }catch (Exception exception){
            log.error("调用glm异常",exception);
        }
        return R.fail("调用异常");
    }

    @Override
    public ResponseBodyEmitter streamChat(StreamChatDTO streamChatDTO, HttpServletResponse response, Consumer<LLMStreamResultDTO> resultHandler) {
        return null;
    }
}
